import os
import cv2
from ultralytics import YOLO
import xml.etree.ElementTree as ET


def draw_bounding_boxes(image, results, class_names):
    """
    在图像上绘制边界框和类别标签
    :param image: 输入的图像
    :param results: 模型推理的结果
    :param class_names: 类别名称列表
    :return: 绘制好边界框的图像
    """
    for result in results:
        boxes = result.boxes
        for box in boxes:
            class_id = int(box.cls.item())
            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(image, class_names[class_id], (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    return image


def save_annotations_xml(image_path, results, class_names, output_xml_folder):
    """
    将标注结果保存为 XML 文件（Pascal VOC 格式）
    :param image_path: 图像的路径
    :param results: 模型推理的结果
    :param class_names: 类别名称列表
    :param output_xml_folder: 输出 XML 文件的文件夹
    """
    image_name = os.path.basename(image_path)
    image = cv2.imread(image_path)
    img_height, img_width = image.shape[:2]
    xml_label_path = os.path.join(output_xml_folder, os.path.splitext(image_name)[0] + '.xml')

    # 创建 XML 根节点
    root = ET.Element('annotation')

    # 添加文件夹和文件名信息
    folder = ET.SubElement(root, 'folder')
    folder.text = os.path.dirname(image_path)
    filename = ET.SubElement(root, 'filename')
    filename.text = image_name
    path = ET.SubElement(root, 'path')
    path.text = image_path

    # 添加图像大小信息
    size = ET.SubElement(root, 'size')
    width_elem = ET.SubElement(size, 'width')
    width_elem.text = str(img_width)
    height_elem = ET.SubElement(size, 'height')
    height_elem.text = str(img_height)
    depth = ET.SubElement(size, 'depth')
    depth.text = '3'

    # 添加标注信息
    for result in results:
        boxes = result.boxes
        for box in boxes:
            class_id = int(box.cls.item())
            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
            obj = ET.SubElement(root, 'object')
            name = ET.SubElement(obj, 'name')
            name.text = class_names[class_id]
            pose = ET.SubElement(obj, 'pose')
            pose.text = 'Unspecified'
            truncated = ET.SubElement(obj, 'truncated')
            truncated.text = '0'
            difficult = ET.SubElement(obj, 'difficult')
            difficult.text = '0'
            bbox = ET.SubElement(obj, 'bndbox')
            xmin = ET.SubElement(bbox, 'xmin')
            xmin.text = str(x1)
            ymin = ET.SubElement(bbox, 'ymin')
            ymin.text = str(y1)
            xmax = ET.SubElement(bbox, 'xmax')
            xmax.text = str(x2)
            ymax = ET.SubElement(bbox, 'ymax')
            ymax.text = str(y2)

    # 保存 XML 文件
    tree = ET.ElementTree(root)
    tree.write(xml_label_path, encoding='utf-8', xml_declaration=True)


def main():
    # 加载本地模型
    model_path = 'best.pt'
    model = YOLO(model_path)

    # 定义类别名称
    class_names = ['Role']  # 根据实际情况修改

    # 输入图片文件夹和输出文件夹
    input_image_folder = './Data/Photo'
    output_image_folder = 'data/data3/JPEGImages'
    output_xml_folder = 'data/data3/Annotations'

    # 创建输出文件夹（如果不存在）
    if not os.path.exists(output_image_folder):
        os.makedirs(output_image_folder)
    if not os.path.exists(output_xml_folder):
        os.makedirs(output_xml_folder)

    # 遍历输入图片文件夹中的所有图片
    for filename in os.listdir(input_image_folder):
        if filename.endswith(('.jpg', '.png')):
            image_path = os.path.join(input_image_folder, filename)
            # 读取图像
            image = cv2.imread(image_path)

            # 使用模型进行推理
            results = model.predict(image)

            # 在图像上绘制边界框
            annotated_image = draw_bounding_boxes(image.copy(), results, class_names)

            # 保存绘制好的图像
            output_image_path = os.path.join(output_image_folder, filename)
            cv2.imwrite(output_image_path, annotated_image)

            # 保存标注结果为 XML 文件
            save_annotations_xml(image_path, results, class_names, output_xml_folder)

            print(f"Processed {filename} and saved results.")


if __name__ == "__main__":
    main()